'''
@Time    : 2022/3/26 12:41
@Author  : Fu Junyu
@Site    : www.fujunyu.cn
@File    : model.py
@Software: PyCharm
'''
import torch
from torch import nn
from HeteroRGCNLayer import HeteroRGCNLayer

'''
构建包含RGCN层的训练模型
'''
class RGCNModel(nn.Module):
    def __init__(self, user_in_size, movies_in_size, out_size, user_max_dict, movie_max_dict, num_bases=3, num_rels=3, regularizer='basis'):
        super(RGCNModel, self).__init__()
        self.user_in_size = user_in_size
        self.movies_in_size = movies_in_size
        self.out_size = out_size
        self.num_bases = num_bases
        self.num_rels = num_rels
        self.regularizer = regularizer
        self.user_max_dict = user_max_dict
        self.movie_max_dict = movie_max_dict

        # --------------------------------- 图卷积模块 ------------------------------------------------------------------
        self.RGCNlayer_1 = HeteroRGCNLayer(self.user_in_size, self.movies_in_size, self.out_size, self.num_bases, self.num_rels, self.regularizer)
        self.RGCNlayer_2 = HeteroRGCNLayer(self.out_size, self.out_size, self.out_size, self.num_bases, self.num_rels, self.regularizer)

        # --------------------------------- 线性全连接层 -----------------------------------------------------------------
        self.linear1 = nn.Linear(out_size, out_size*2)
        self.linear2 = nn.Linear(out_size, out_size*2)

        self.linear1_BN = nn.BatchNorm1d(out_size*2)
        self.linear2_BN = nn.BatchNorm1d(out_size*2)

        self.act = nn.Sigmoid()

        self.projectionMatrix = nn.Linear(out_size*2, out_size*2)

        self.Br1 = nn.Parameter(torch.Tensor(out_size*2, 1))
        nn.init.xavier_uniform_(self.Br1, gain=nn.init.calculate_gain('sigmoid'))
        self.Br2 = nn.Parameter(torch.Tensor(out_size*2, 1))
        nn.init.xavier_uniform_(self.Br2, gain=nn.init.calculate_gain('sigmoid'))

        self.Ye_BN = nn.BatchNorm1d(1)

        self.final = nn.Linear(out_size*4, 5)


    def forward(self, graph, feat_dic, trainData):
        uid =feat_dic['user']['uid']
        gender = feat_dic['user']['gender']
        age = feat_dic['user']['age']
        job = feat_dic['user']['job']
        zip = feat_dic['user']['zip']
        feat_user = torch.cat((uid.view(uid.shape[0], 1), gender.view(gender.shape[0], 1)), 1)
        feat_user = torch.cat((feat_user, age.view(age.shape[0], 1)), 1)
        feat_user = torch.cat((feat_user, job.view(job.shape[0], 1)), 1)
        feat_user = torch.cat((feat_user, zip.view(zip.shape[0], 1)), 1)

        mid = feat_dic['movies']['mid']
        mtype = feat_dic['movies']['mtype']
        mtext = feat_dic['movies']['mtext']
        feat_movies = torch.cat((mid.view(mid.shape[0], 1), mtype), 1)
        feat_movies = torch.cat((feat_movies, mtext), 1)

        feat_dic = {'user': feat_user.float(), 'movies': feat_movies}

        user_movies_feat = self.RGCNlayer_1(graph, feat_dic)
        user_movies_feat = self.RGCNlayer_2(graph, user_movies_feat)

        rgcn_feat_user = user_movies_feat['user']
        rgcn_feat_movies = user_movies_feat['movies']

        trainUser = torch.index_select(rgcn_feat_user, 0, trainData['userID'])
        trainMovies = torch.index_select(rgcn_feat_movies, 0, trainData['moviesID'])

        trainUser = self.linear1(trainUser)
        trainMovies = self.linear2(trainMovies)

        Ye1 = self.act(self.projectionMatrix(trainUser))
        Ye2 = self.act(self.projectionMatrix(trainMovies))

        Ye1 = torch.matmul(Ye1, self.Br1)
        Ye2 = torch.matmul(Ye2, self.Br2)

        label_pred = Ye1 * Ye2
        label_pred = label_pred.sum(1)


        return label_pred




